from enum import Enum
from typing import Dict, Callable

from pyspark import Row
from pyspark.ml.param import Param, Params
from pyspark.sql import DataFrame

from gai.v2.spark.base import SparkTransformer
from gai.v2.spark.demo_data import demo_data_for_missing_value_handler
from gai.v2.utils import get_or_create_spark_session


class MissingValueStrategy(Enum):
    """An enumeration of action strategies for the case when a cell is found null.
    """
    DELETE_MISSING = 1
    "delete the row when the specified cell is null"
    SUBSTITUTE_MISSING = 2
    "default, substitute the cell value"


class SubstitutionValue(Enum):
    """An enumeration of substitution value strategies for the case when a cell
    is found null.
    """

    FILL_MINIMUM = 1
    "fill with the column minimum when the specified cell is null"
    FILL_MAXIMUM = 2
    "fill with the column maximum when the specified cell is null"
    FILL_MEAN = 3
    "fill with the column mean when the specified cell is null"


def handle_missing_values(dataset: DataFrame, mval_strategy: Dict, subst_values: Dict):
    """Given a dataframe, a group of missing value strategies, and a group of
    substitution values, handle the missing values as specified by the strategies
    and the substitution values.

    Args:
        dataset:
            the dataframe.
        mval_strategy:
            a dictionary specifying the missing value strategies.
        subst_values:
            a dictionary specifying the substitution values which can be either
            a concrete value or an abstract value (that is, substitution
            strategy).

    Returns:
        the dataframe with missing values dealt with.

    See Also:
        :class:`MissingValueHandler`
    """

    # preliminary check
    column_set = set(dataset.columns)
    assert set(mval_strategy.keys()) <= column_set
    assert set(subst_values.keys()) <= column_set

    # carry out the missing value handling
    filter_pred = make_filter_predicate(mval_strategy)
    materialized = materialize_substitution_values(dataset, subst_values)
    partition_mapper = make_partition_mapper(filter_pred, materialized)

    ret = dataset.rdd.mapPartitions(partition_mapper)
    spark = get_or_create_spark_session()
    return spark.createDataFrame(ret, dataset.schema)


def make_partition_mapper(filter_pred: Callable, subst_values: Dict):
    """Given a filter predicate and a dictionary of substitution values, return
    a partition mapper function that filters out all the rows that fail to
    satisfy the predicate, and replaces the null cells of the remaining rows
    with the specified substitution values. If some cell is null but the
    substitution value for its column is unspecified, it is preserved as a null
    cell.

    Args:
        filter_pred:
            the filter predicate.
        subst_values:
            a dictionary specifying the substitution values for certain columns.
            Its form is like ``{'col1': 'val1', 'col2': 'val2'}``.

    Returns:
        the desired partition mapper.
    """

    def partition_mapper(rows):
        for row in rows:
            if not filter_pred(row):
                continue

            d = row.asDict()
            modified = False
            for k, v in subst_values.items():
                if d[k] is None:
                    d[k] = v
                    modified = True

            if modified:
                yield Row(**d)
            else:
                yield row

    return partition_mapper


def make_filter_predicate(mval_strategy: Dict):
    """Given a missing-value strategy, return a predicate on ``Row`` that is
    false when any of the columns bound to ``DELETE_MISSING`` is null.

    Args:
        mval_strategy:
            the missing-value strategy.

    Returns:
        the predicate.

    >>> f = make_filter_predicate({
    ...      "gid_": MissingValueStrategy.DELETE_MISSING,
    ...      "ft_category_cnt": MissingValueStrategy.SUBSTITUTE_MISSING
    ...     })
    >>> f(Row(gid_=None, ft_usertags="", ft_category_cnt=None))
    False
    >>> f(Row(gid_="some_value", ft_usertags="", ft_category_cnt=None))
    True
    """

    deletion_set = [name for name, strategy in mval_strategy.items()
                    if strategy == MissingValueStrategy.DELETE_MISSING]

    def filter_pred(row: Row):
        return all(row[col] is not None for col in deletion_set)

    return filter_pred


def materialize_substitution_values(dataset: DataFrame, subst_values: Dict):
    """Given a dataset and a dictionary specifying the substitution values,
    materialize the abstract substitution strategies with concrete values.

    Consider any column. If its substitution strategies is ``FILL_MINIMUM``,
    the materialized value is the minimum value for the column. Other strategies
    (``FILL_MAXIMUM`` and ``FILL_MEAN``) are dealt with similarly. If the
    substitution value is already a concrete value, then this value is used.

    Args:
        dataset:
            the data frame.
        subst_values:
            the substitution values.
    Returns:
        the materialized substitution values.


    >>> df = demo_data_for_missing_value_handler()
    >>> df.show()
    +------+--------+--------------------+-------------+--------------------+---------------+
    |secret|     day|                gid_|_month_offset|         ft_usertags|ft_category_cnt|
    +------+--------+--------------------+-------------+--------------------+---------------+
    |     1|20181114|ANDROID-68b6c301c...|         null|                    |           null|
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|            0|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     3|20181114|                null|            0|                    |           null|
    |     3|20181114|                null|           -1|                    |           null|
    |     4|20181126|ANDROID-0000db9c0...|            0|                    |              3|
    |     4|20181126|ANDROID-0000db9c0...|           -1|                    |              3|
    |     5|20181126|ANDROID-846049d18...|            0|022000,026400,02d400|             12|
    |     5|20181126|ANDROID-846049d18...|           -1|022000,026400,02d400|             13|
    |     7|20181126|                null|            0|                    |           null|
    |     7|20181126|                null|           -1|                    |           null|
    |     8|20180709|ANDROID-000340aa5...|            0|    02b000,02b100,h0|             14|
    |     8|20180709|ANDROID-000340aa5...|           -1|    02b000,02b100,h0|             15|
    |     9|20180709|ANDROID-00033a1c3...|            0|022000,022500,02d300|              7|
    |     9|20180709|ANDROID-00033a1c3...|           -1|022000,022100,02d300|              9|
    |    11|20180709|                null|            0|                    |           null|
    |    11|20180709|                null|           -1|                    |           null|
    +------+--------+--------------------+-------------+--------------------+---------------+
    <BLANKLINE>
    >>> subst_values = {
    ...    "_month_offset"  : SubstitutionValue.FILL_MINIMUM,
    ...    "ft_category_cnt": SubstitutionValue.FILL_MEAN,
    ...    "gid_"           : "non-existent"
    ... }
    >>> materialized = materialize_substitution_values(df, subst_values)
    >>> print(materialized)
    {'_month_offset': -1, 'ft_category_cnt': 9.5, 'gid_': 'non-existent'}
    """

    ret = subst_values.copy()

    import pyspark.sql.functions as F
    for col, value in subst_values.items():
        if value == SubstitutionValue.FILL_MAXIMUM:
            ret[col] = dataset.select(F.max(col)).first()['max(' + col + ')']
        elif value == SubstitutionValue.FILL_MINIMUM:
            ret[col] = dataset.select(F.min(col)).first()['min(' + col + ')']
        elif value == SubstitutionValue.FILL_MEAN:
            ret[col] = dataset.select(F.mean(col)).first()['avg(' + col + ')']

    return ret


class MissingValueHandler(SparkTransformer):
    """Given a dataframe, a group of missing value strategies, and a group of
    substitution values, handle the missing values as specified by the strategies
    and the substitution values.

    Args:
        mvalStrategy:
            a dictionary specifying the missing value strategies.
        substValues:
            a dictionary specifying the substitution values which can be either
            a concrete value or an abstract value (that is, substitution
            strategy).



    >>> df = demo_data_for_missing_value_handler()
    >>> df.show()
    +------+--------+--------------------+-------------+--------------------+---------------+
    |secret|     day|                gid_|_month_offset|         ft_usertags|ft_category_cnt|
    +------+--------+--------------------+-------------+--------------------+---------------+
    |     1|20181114|ANDROID-68b6c301c...|         null|                    |           null|
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|            0|                    |           null|
    |     2|20181114|ANDROID-68b6c301c...|           -1|                    |           null|
    |     3|20181114|                null|            0|                    |           null|
    |     3|20181114|                null|           -1|                    |           null|
    |     4|20181126|ANDROID-0000db9c0...|            0|                    |              3|
    |     4|20181126|ANDROID-0000db9c0...|           -1|                    |              3|
    |     5|20181126|ANDROID-846049d18...|            0|022000,026400,02d400|             12|
    |     5|20181126|ANDROID-846049d18...|           -1|022000,026400,02d400|             13|
    |     7|20181126|                null|            0|                    |           null|
    |     7|20181126|                null|           -1|                    |           null|
    |     8|20180709|ANDROID-000340aa5...|            0|    02b000,02b100,h0|             14|
    |     8|20180709|ANDROID-000340aa5...|           -1|    02b000,02b100,h0|             15|
    |     9|20180709|ANDROID-00033a1c3...|            0|022000,022500,02d300|              7|
    |     9|20180709|ANDROID-00033a1c3...|           -1|022000,022100,02d300|              9|
    |    11|20180709|                null|            0|                    |           null|
    |    11|20180709|                null|           -1|                    |           null|
    +------+--------+--------------------+-------------+--------------------+---------------+
    <BLANKLINE>
    >>> mval_strategy = {
    ...     "gid_": MissingValueStrategy.DELETE_MISSING
    ... }
    >>> subst_values = {
    ...     "_month_offset"  : SubstitutionValue.FILL_MINIMUM,
    ...     "ft_category_cnt": SubstitutionValue.FILL_MEAN,
    ... }
    >>> handler = MissingValueHandler(mval_strategy, subst_values)
    >>> result = handler.transform(df)
    >>> same_result = handle_missing_values(df, mval_strategy, subst_values)
    >>> result.show()
    +------+--------+--------------------+-------------+--------------------+---------------+
    |secret|     day|                gid_|_month_offset|         ft_usertags|ft_category_cnt|
    +------+--------+--------------------+-------------+--------------------+---------------+
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |            9.5|
    |     1|20181114|ANDROID-68b6c301c...|           -1|                    |            9.5|
    |     2|20181114|ANDROID-68b6c301c...|            0|                    |            9.5|
    |     2|20181114|ANDROID-68b6c301c...|           -1|                    |            9.5|
    |     4|20181126|ANDROID-0000db9c0...|            0|                    |              3|
    |     4|20181126|ANDROID-0000db9c0...|           -1|                    |              3|
    |     5|20181126|ANDROID-846049d18...|            0|022000,026400,02d400|             12|
    |     5|20181126|ANDROID-846049d18...|           -1|022000,026400,02d400|             13|
    |     8|20180709|ANDROID-000340aa5...|            0|    02b000,02b100,h0|             14|
    |     8|20180709|ANDROID-000340aa5...|           -1|    02b000,02b100,h0|             15|
    |     9|20180709|ANDROID-00033a1c3...|            0|022000,022500,02d300|              7|
    |     9|20180709|ANDROID-00033a1c3...|           -1|022000,022100,02d300|              9|
    +------+--------+--------------------+-------------+--------------------+---------------+
    <BLANKLINE>
    """

    mvalStrategy = Param(Params._dummy(), "mvalStrategy",
                         "a dictionary specifying whether to keep or remove the "
                         "record if a cell of specified column is found null")
    substValues = Param(Params._dummy(), "substValues",
                        "a dictionary specifying what kind of values to "
                        "substitute the cell if a cell of specified column is "
                        "found null")

    def __init__(self, mvalStrategy: Dict, substValues: Dict):
        super(MissingValueHandler, self).__init__()

        self.setMvalStrategy(mvalStrategy) \
            .setSubstValues(substValues)

    def setMvalStrategy(self, mvalStrategy):
        self._paramMap[self.mvalStrategy] = mvalStrategy
        return self

    def getMvalStrategy(self):
        return self.getOrDefault(self.mvalStrategy)

    def setSubstValues(self, substValues):
        self._paramMap[self.substValues] = substValues
        return self

    def getSubstValues(self):
        return self.getOrDefault(self.substValues)

    def _transform(self, dataset):
        return handle_missing_values(dataset, self.getMvalStrategy(), self.getSubstValues())
